Demonstrates autograd integration with NVFuser multidevice #3787
Conversation
|
Cool -- add me to reviewers when it's ready! |
|
@wujingyue To review. |
|
Oops I can't add to the reviewers list. |
tests/python/test_dtensor.py
Outdated
| self.s = sequence | ||
| self.e = hidden | ||
|
|
||
| class LinearForwardDefinition(FusionDefintionArguments): |
There was a problem hiding this comment.
I feel using class and inheritance is an overkill. Functions and partials should be good enough.
def define_linear_forward(config: LinearConfig, fd: FusionDefinition) -> None:
and later
partial(define_linear_forward, config)
tests/python/test_dtensor.py
Outdated
| ): | ||
| b, s, e = input._local_tensor.shape | ||
| d = weight.device_mesh.size() | ||
| op = FusionDefinitionWrapper(LinearForwardDefinition(d, b, s, e)) |
There was a problem hiding this comment.
Can you try to construct the op in __init__? Example: https://github.com/canqin001/PointDAN/blob/5001b38cb5506b1c6b40ad1329c1d6f4fbbdd26d/Model.py#L29. I'm worried about the overhead of constructing FusionDefinitionWrapper for each forward and backward call.
There was a problem hiding this comment.
I don't think we can create init for torch.autograd.Function: https://github.com/pytorch/pytorch/blob/main/torch/autograd/function.py#L498-L506. Also looks like the code from the link is quite old. I get the following error when trying to add __init__ and use .apply:
# When calling, LinearFunction(d, b, s, e).apply(inp_dtensor, weight_dtensor, bias_dtensor)
[rank0]: Traceback (most recent call last):
[rank0]: File "/opt/home/dtensor_extension/test.py", line 207, in <module>
[rank0]: out_dtensor = LinearFunction(d, b, s, e).apply(inp_dtensor, weight_dtensor, bias_dtensor)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 575, in apply
[rank0]: return super().apply(*args, **kwargs) # type: ignore[misc]
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/home/dtensor_extension/test.py", line 161, in forward
[rank0]: outputs = self.forward_op([input, weight, bias])
[rank0]: ^^^^
[rank0]: NameError: name 'self' is not defined
There was a problem hiding this comment.
Thanks for confirming! We don't need to figure this out for this PR, but I still wonder how to avoid creating a FusionDefinitionWrapper for each call.
|
!test |
|
!test |
|
@wujingyue Addressed the review. |
|
There are some real errors in nvfuser-ci/jit_python_distributed_tests_20_A100. |
|
!test |
|
@wujingyue This is ready to be merged. |
This PR demonstrates how to wrap a forward and a backward fusion definition in a
torch.autograd.Functionthat takes PyTorch DTensors as input and outputs PyTorch DTensors.